#!/usr/bin/env python
# -*- coding: utf-8 -*-
from __future__ import print_function
from __future__ import absolute_import
import sys
import subprocess

from db import Base, Session
from db.models import User12306cn, UserCSDN, UserDuoWan, UserDuDuNiu,\
        User17173com, UserTianYa
from db.models import UserWangYi, UserMao, UserShengDa, User178com, User7k7k,\
        User163com, User163mail

def num_of_lines(filename):
    """fast to get number of lines for large file"""
    p = subprocess.Popen(['wc', '-l', filename],
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE)
    result, err = p.communicate()
    if p.returncode != 0:
        raise IOError(err)
    return int(result.strip().split()[0])

class BaseLoader(object):
    target = None
    delimiter = '\t'
    def __init__(self):
        self.fin = sys.stdin
        self.fout = sys.stdout
        self.ferr = sys.stderr

    def load(self, session, filename, append=True, bulk=1000):
        if not self.handle_begin(filename=filename):
            return False
        self.target.__table__.create(session.get_bind(), checkfirst=True)
        print('[+] {} file {} into TABLE `{}`'.format(
            'Appending' if append else '(Re)Loading',
            filename, self.target.__tablename__))
        if not append:
            rows_deleted = session.query(self.target).delete()
            print('[+] {} rows deleted from existing table'.format(rows_deleted))
        success_count = 0
        fail_count = 0
        total_count = num_of_lines(filename)
        print('[+] Total {} lines to be read'.format(total_count))
        for line in self.fin:
            if self.handle_line(session, line):
                success_count += 1
                self.fout.write(line)
                # better idea?
                if success_count % bulk == 0:
                    session.commit()
                    sys.stdout.write('\r[+] Reading {:.2f}% ({} failed)'.format(
                        (success_count+fail_count)*100.0/(total_count),
                        fail_count))
                    sys.stdout.flush()
            else:
                fail_count += 1
                self.ferr.write(line)
        print()
        self.handle_end()
        session.commit()
        print('[+] Done. {} success, {} failed'.format(
            success_count, fail_count))

    def handle_begin(self, **kwargs):
        """used by subclass"""
        filename = kwargs.get('filename', '')
        append = kwargs.get('append', False)
        if filename:
            success = filename + '.success'
            review = filename + '.review'
            self.fin = open(filename, 'r')
            self.fout = open(success, 'w')
            self.ferr = open(review, 'w')
        return True

    def handle_end(self):
        """used by subclass"""
        if self.fin is not sys.stdin:
            self.fin.close()
        if self.ferr is not sys.stderr:
            self.ferr.close()

    def handle_line(self, session, line):
        """override by subclass"""
        data = line.split(self.delimiter)
        if len(data) != 2:
            return False
        try:
            account = data[0].decode('utf-8')
            password = data[1].rstrip().decode('utf-8')
        except UnicodeDecodeError:
            return False
        if len(account) < 3 or len(password) < 6:
            return False
        user = self.target(account=account, password=password)
        session.add(user)
        return True

class Loader12306(BaseLoader):
    target = User12306cn
    def __init__(self):
        super(Loader12306, self).__init__()
    def handle_line(self, session, line):
        data = line.split('----')
        if len(data) != 7:
            return False
        user = self.target(account=data[0],
            password=data[1],
            name=data[2].decode('utf-8'),
            idnum=data[3],
            nickname=data[4],
            phone=data[5])
        session.add(user)
        return True

class LoaderCSDN(BaseLoader):
    target = UserCSDN
    def __init__(self):
        super(LoaderCSDN, self).__init__()
    def handle_line(self, session, line):
        data = line.split(' # ')
        if len(data) != 3:
            return False
        try:
            account = data[0].decode('utf-8')
            password = data[1].decode('utf-8')
            email = data[2].rstrip().decode('utf-8')
        except UnicodeDecodeError:
            return False
        user = self.target(account=account,
                password=password,
                email=email)
        session.add(user)
        return True

class LoaderDuoWan(BaseLoader):
    target = UserDuoWan
    def __init__(self):
        super(LoaderDuoWan, self).__init__()
    def handle_line(self, session, line):
        data = line.split('\t')
        if len(data) != 5:
            return False
        try:
            nickname = data[1].decode('utf-8')
            password = data[2].decode('utf-8')
            email = data[3].decode('utf-8')
            account = data[4].rstrip().decode('utf-8')
        except UnicodeDecodeError:
            return False
        user = self.target(account=account,
                password=password,
                nickname=nickname,
                email=email)
        session.add(user)
        return True

class LoaderDuDuNiu(BaseLoader):
    target = UserDuDuNiu
    def __init__(self):
        super(LoaderDuDuNiu, self).__init__()
    def handle_line(self, session, line):
        data = line.split('----')
        if len(data) != 3:
            return False
        try:
            account = data[0].decode('utf-8')
            email = data[1].decode('utf-8')
            password = data[2].rstrip().decode('utf-8')
        except UnicodeDecodeError:
            return False
        user = self.target(account=account,
                email=email, password=password)
        session.add(user)
        return True

class Loader17173com(BaseLoader):
    target = User17173com
    def __init__(self):
        super(Loader17173com, self).__init__()
    def handle_line(self, session, line):
        data = line.split('----')
        if len(data) != 4:
            return False
        try:
            account = data[0].decode('utf-8')
            password_md5 = data[1].decode('utf-8')
            email = data[2].decode('utf-8')
            password = data[3].rstrip().decode('utf-8')
        except UnicodeDecodeError:
            return False
        user = self.target(account=account, password=password,
                password_md5=password_md5, email=email)
        session.add(user)
        return True

class LoaderTianYa(BaseLoader):
    target = UserTianYa
    def __init__(self):
        super(LoaderTianYa, self).__init__()
    def handle_line(self, session, line):
        data = line.split('\t')
        if len(data) != 3:
            return False
        try:
            account = data[0].decode('utf-8')
            password = data[1].decode('utf-8')
            email = data[2].decode('utf-8')
        except UnicodeDecodeError:
            return False
        user = self.target(account=account,
                password=password, email=email)
        session.add(user)
        return True

class LoaderWangYi(BaseLoader):
    target = UserWangYi
    delimiter = '----'
    def __init__(self):
        super(LoaderWangYi, self).__init__()

class LoaderMao(BaseLoader):
    target = UserMao
    delimiter = '\t'
    def __init__(self):
        super(LoaderMao, self).__init__()

class LoaderShengDa(BaseLoader):
    target = UserShengDa
    delimiter = '----'
    def __init__(self):
        super(LoaderShengDa, self).__init__()

class Loader178com(BaseLoader):
    target = User178com
    delimiter = '----'
    def __init__(self):
        super(Loader178com, self).__init__()

class Loader7k7k(BaseLoader):
    target = User7k7k
    delimiter = '----'
    def __init__(self):
        super(Loader7k7k, self).__init__()

class Loader163com(BaseLoader):
    target = User163com
    delimiter = '----'
    def __init__(self):
        super(Loader163com, self).__init__()

class Loader163mail(BaseLoader):
    target = User163mail
    delimiter = '----'
    def __init__(self):
        super(Loader163mail, self).__init__()
